{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1df2a482",
   "metadata": {},
   "source": [
    "# Joint Differential Equations\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_toolbox/joint_equations.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_toolbox/joint_equations.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "109f9b4e",
   "metadata": {},
   "source": [
    "@[Xiaoyu Chen](mailto:c-xy17@tsinghua.org.cn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9df7780",
   "metadata": {},
   "source": [
    "In a [dynamical system](../tutorial_building/dynamical_systems.ipynb), there may be multiple variables that change dynamically over time. Sometimes these variables are interconnected, and updating one variable requires others as the input. For example, in the widely known Hodgkin–Huxley model, the variables $V$, $m$, $h$, and $n$ are updated synchronously and interdependently (please refer to [Building Neuron Models](../tutorial_building/neuron_models.ipynb)for details). To achieve higher integral accuracy, it is recommended to use ``brainpy.JointEq`` to jointly solving interconnected differential equations."
   ]
  },
  {
   "cell_type": "code",
   "id": "be08d171",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.420059Z",
     "start_time": "2025-10-06T05:15:55.413203Z"
    }
   },
   "source": "import brainpy as bp",
   "outputs": [],
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "id": "991cf807",
   "metadata": {},
   "source": [
    "## ``brainpy.JointEq``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05d270a9",
   "metadata": {},
   "source": [
    "``brainpy.JointEq`` is used to merge individual but interconnected differential equations into a single joint equation. For example, below are the two differential equations of the Izhikevich model:"
   ]
  },
  {
   "cell_type": "code",
   "id": "2921b856",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.429070Z",
     "start_time": "2025-10-06T05:15:55.423341Z"
    }
   },
   "source": [
    "a, b = 0.02, 0.20\n",
    "dV = lambda V, t, u, Iext: 0.04 * V * V + 5 * V + 140 - u + Iext\n",
    "du = lambda u, t, V: a * (b * V - u)"
   ],
   "outputs": [],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "id": "44527a26",
   "metadata": {},
   "source": [
    "Where updating $V$ requires $u$ as the input, and updating $u$ requires $V$ as the input. The joint equation can be defined as:"
   ]
  },
  {
   "cell_type": "code",
   "id": "08ac3b75",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.435687Z",
     "start_time": "2025-10-06T05:15:55.432079Z"
    }
   },
   "source": [
    "joint_eq = bp.JointEq(dV, du)"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "id": "7dcfcc88",
   "metadata": {},
   "source": [
    "``brainpy.JointEq`` receives only one argument named `eqs`, which can be a list or tuple containing multiple differential equations. Then it can be packed into a numarical integrator that solves the equation with a specified method, just as what can be done to any individual differential equation."
   ]
  },
  {
   "cell_type": "code",
   "id": "356cf60d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.443030Z",
     "start_time": "2025-10-06T05:15:55.439434Z"
    }
   },
   "source": [
    "itg = bp.odeint(joint_eq, method='rk2')"
   ],
   "outputs": [],
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "id": "1145f933",
   "metadata": {},
   "source": [
    "There are several requirements for defining a joint equation:\n",
    "1. Every individual differential equation should follow the format of defining a [ODE](ode_numerical_solvers.ipynb) or [SDE](sde_numerical_solvers.ipynb) funtion in BrainPy. For example, the arguments before `t` denote the dynamical variables and arguments after `t` denote the parameters.\n",
    "2. The same variable in different equations should have the same name. Different variables should named differently.\n",
    "\n",
    "Note that `brainpy.JointEq` supports make nested ``JointEq``, which means the instance of ``JointEq`` can be an element to compose a new ``JointEq``."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f7c666",
   "metadata": {},
   "source": [
    "## Why use `brainpy.JointEq`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3996db1",
   "metadata": {},
   "source": [
    "Users may be confused with the function of `brainpy.JointEq`, because multiple differential equations can be written in a single function:"
   ]
  },
  {
   "cell_type": "code",
   "id": "4dec7537",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.450726Z",
     "start_time": "2025-10-06T05:15:55.447321Z"
    }
   },
   "source": [
    "def diff(V, u, t, Iext):\n",
    "    dV = 0.04 * V * V + 5 * V + 140 - u + Iext\n",
    "    du = a * (b * V - u)\n",
    "    return dV, du\n",
    "\n",
    "itg_V_u = bp.odeint(diff, method='rk2')"
   ],
   "outputs": [],
   "execution_count": 13
  },
  {
   "cell_type": "markdown",
   "id": "5943bc7a",
   "metadata": {},
   "source": [
    "or simply packed into interators separately:"
   ]
  },
  {
   "cell_type": "code",
   "id": "12e5d88d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.460116Z",
     "start_time": "2025-10-06T05:15:55.455524Z"
    }
   },
   "source": [
    "int_V = bp.odeint(dV, method='rk2')\n",
    "int_u = bp.odeint(du, method='rk2')"
   ],
   "outputs": [],
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "id": "50963fe1",
   "metadata": {},
   "source": [
    "To illusrate the difference between joint and separate differential equations, let's dive into the differential codes of these two types of equations. \n",
    "\n",
    "If we make numerical solver for each derivative function, they will be solved independently:"
   ]
  },
  {
   "cell_type": "code",
   "id": "38101bec",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.473617Z",
     "start_time": "2025-10-06T05:15:55.467195Z"
    }
   },
   "source": "bp.odeint(dV, method='rk2', show_code=True)",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def brainpy_itg_of_ode10(V, t, u, Iext, dt=0.1):\n",
      "  dV_k1 = f(V, t, u, Iext)\n",
      "  k2_V_arg = V + dt * dV_k1 * 0.6666666666666666\n",
      "  k2_t_arg = t + dt * 0.6666666666666666\n",
      "  dV_k2 = f(k2_V_arg, k2_t_arg, u, Iext)\n",
      "  V_new = V + dV_k1 * dt * 0.25 + dV_k2 * dt * 0.75\n",
      "  return V_new\n",
      "\n",
      "{'f': <function <lambda> at 0x0000015BBFC96C00>}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<brainpy.integrators.ode.explicit_rk.RK2 at 0x15bbfbdd040>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "id": "8500a6c7",
   "metadata": {},
   "source": [
    "As is shown in the output code, the variable $V$ is integrated twice by the RK2 method. For the second differential value `dV_k2`, the updated value of $V$ (`k2_V_arg`) and original $u$ are used to calculate the differential value. This will generate a tiny error, since the values of $V$ and $u$ are taken at different times.\n",
    "\n",
    "To eliminate this error, the differential equation of $V$ and $u$ should be solved jointly through `brainpy.JointEq`:"
   ]
  },
  {
   "cell_type": "code",
   "id": "32901ae6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:15:55.494272Z",
     "start_time": "2025-10-06T05:15:55.488118Z"
    }
   },
   "source": [
    "eq = bp.JointEq(dV, du)\n",
    "bp.odeint(eq, method='rk2', show_code=True)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def brainpy_itg_of_ode11_joint_eq(V, u, t, Iext, dt=0.1):\n",
      "  dV_k1, du_k1 = f(V, u, t, Iext)\n",
      "  k2_V_arg = V + dt * dV_k1 * 0.6666666666666666\n",
      "  k2_u_arg = u + dt * du_k1 * 0.6666666666666666\n",
      "  k2_t_arg = t + dt * 0.6666666666666666\n",
      "  dV_k2, du_k2 = f(k2_V_arg, k2_u_arg, k2_t_arg, Iext)\n",
      "  V_new = V + dV_k1 * dt * 0.25 + dV_k2 * dt * 0.75\n",
      "  u_new = u + du_k1 * dt * 0.25 + du_k2 * dt * 0.75\n",
      "  return V_new, u_new\n",
      "\n",
      "{'f': <brainpy.integrators.joint_eq.JointEq object at 0x0000015BBFBDD7F0>}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<brainpy.integrators.ode.explicit_rk.RK2 at 0x15bbfbdd790>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "id": "6a1a9669",
   "metadata": {},
   "source": [
    "It is shown in this output code that second differential values of $v$ and $u$ are calculated by using the updated values (`k2_V_arg` and `k2_u_arg`) at the same time. This will result in a more accurate integral."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73051bec",
   "metadata": {},
   "source": [
    "The figure below compares the simulation results of the Izhikevich model using joint and separate differential equations ($dt = 0.2 ms$). It is shown that as the simulation time increases, the integral error becomes greater.\n",
    "\n",
    "<img src=\"../_static/joint_and_separate_equations.png\" width=\"900 px\">"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
